package fiveagon.lsystem;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;
import java.util.Random;

public class StochasticGrammar extends AbstractGrammar {
  private final Map<Character, Production> ruleMap = new HashMap<>();

  public StochasticGrammar(String axiom) {
    super(axiom);
  }

  public void addRule(char c, String s) {
    addRule(c, s, 1.0f);
  }

  public void addRule(char c, String s, float weight) {
    Character key = Character.valueOf(c);
    if (ruleMap.containsKey(key)) {
      ruleMap.get(key).addProduction(s, weight);
    } else {
      ruleMap.put(key, new Production(s, weight));
    }
  }

  @Override
  String generate(String s) {
    StringBuilder stringBuilder = new StringBuilder();
    for (int i = 0, n = s.length(); i < n; i++) {
      char c = s.charAt(i);
      Character key = Character.valueOf(c);
      if (ruleMap.containsKey(key)) {
        stringBuilder.append(ruleMap.get(key).getProduction());
      } else {
        stringBuilder.append(c);
      }
    }
    return stringBuilder.toString();
  }

  private static class Production {
    private static final Random RANDOM = new Random();
    private final Collection<WeightedProduction> productionCollection = new ArrayList<>();
    private float totalWeight = 0.0f;

    private Production(String s, float weight) {
      addProduction(s, weight);
    }

    private void addProduction(String s, float weight) {
      productionCollection.add(new WeightedProduction(s, weight));
      totalWeight += weight;
    }

    private String getProduction() {
      float randomPick = RANDOM.nextFloat() * totalWeight;
      float cumulative = 0.0f;
      for (WeightedProduction wp : productionCollection) {
        cumulative += wp.weight;
        if (cumulative >= randomPick) {
          return wp.production;
        }
      }
      throw new RuntimeException("Random rule chooser failure");
    }

    private static class WeightedProduction {
      private final String production;
      private final float weight;

      private WeightedProduction(String production, float weight) {
        this.production = production;
        this.weight = weight;
      }
    }
  }
}
